import os
import argparse
import sys

sys.path.append('../')  # 确保这个路径是正确的
from cyw_devkit.core.dataset import CywDataset
from my_io import copy_file
from tictoc import TicToc
import multiprocessing


def main(args):
    # 统计耗时
    cost = TicToc("标注数据生成")
    assert os.path.exists(args.data_path)
    frames = os.listdir(args.data_path)
    frames.sort(key=lambda x: x)
    files = []
    for dir in frames:
        if dir[:2] == '__':  # '__'
            dir = os.path.join(args.data_path, dir)
            if os.path.isdir(dir):
                files.append(dir)

    process_size = len(files)
    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


def copy_files(dir, dst):
    files = os.listdir(dir)
    for file in files:
        copy_file(os.path.join(dir, file), dst)


def main_worker(dataset_path, args):
    lidar_path = os.path.join(dataset_path, 'lidar_bin')
    copy_files(lidar_path, os.path.join(args.anno_path, 'lidar'))

    label_path = os.path.join(dataset_path, 'label')
    copy_files(label_path, os.path.join(args.anno_path, 'label'))

    image_path = os.path.join(dataset_path, 'camera/camera75')
    copy_files(image_path, os.path.join(args.anno_path, 'camera/camera75'))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--data-path', default="/home/adt/bags/work_space/datasets",
                        help='your data root')
    parser.add_argument('--anno-path', default="/home/adt/docker_space/test_carla",
                        help='')
    args = parser.parse_args()

    main(args)
